{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "view-in-github"
   },
   "source": [
    "<a href=\"https://colab.research.google.com/github/google/neural-tangents/blob/main/notebooks/empirical_ntk_resnet.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "nTt0UNQbk_Td"
   },
   "source": [
    "# Example of computing NTK of a ResNet18 on ImageNet inputs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "dvXMUSdFjCqq"
   },
   "source": [
    "Tested on NVIDIA V100"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Sl2JyRE1hK-z"
   },
   "source": [
    "# Imports and setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "HunkuGjSr63O",
    "outputId": "a4d4fda7-d6ec-46ef-9201-b599c44b8714",
    "colab": {
     "base_uri": "https://localhost:8080/"
    }
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "GPU 0: Tesla V100-SXM2-16GB (UUID: GPU-de050d5b-0f8c-274e-ced2-c31bb2d7be5c)\n"
     ]
    }
   ],
   "source": [
    "!nvidia-smi -L"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "bmmbCtfh76RS"
   },
   "outputs": [],
   "source": [
    "!pip install -q --upgrade pip\n",
    "!pip install -q jax[cuda11_cudnn82] -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html\n",
    "!pip install -q git+https://github.com/google/flax\n",
    "!pip install -q git+https://www.github.com/google/neural-tangents"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "4gkWtXVMD1MT"
   },
   "outputs": [],
   "source": [
    "from functools import partial\n",
    "from typing import Any, Callable, Sequence, Optional\n",
    "from flax import linen as nn\n",
    "from jax import jit\n",
    "from jax import numpy as jnp\n",
    "from jax import random\n",
    "\n",
    "import neural_tangents as nt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "dNMGKyttrRkC"
   },
   "source": [
    "# ResNet18 definition, copied from [FLAX examples](https://github.com/google/flax/blob/main/examples/imagenet/models.py)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "3HG2K1ls5ynG"
   },
   "outputs": [],
   "source": [
    "_ModuleDef = Any\n",
    "\n",
    "\n",
    "class _ResNetBlock(nn.Module):\n",
    "  \"\"\"ResNet block.\"\"\"\n",
    "  filters: int\n",
    "  conv: _ModuleDef\n",
    "  norm: _ModuleDef\n",
    "  act: Callable\n",
    "  strides: tuple[int, int] = (1, 1)\n",
    "\n",
    "  @nn.compact\n",
    "  def __call__(self, x,):\n",
    "    residual = x\n",
    "    y = self.conv(self.filters, (3, 3), self.strides)(x)\n",
    "    y = self.norm()(y)\n",
    "    y = self.act(y)\n",
    "    y = self.conv(self.filters, (3, 3))(y)\n",
    "    y = self.norm(scale_init=nn.initializers.zeros)(y)\n",
    "\n",
    "    if residual.shape != y.shape:\n",
    "      residual = self.conv(self.filters, (1, 1),\n",
    "                           self.strides, name='conv_proj')(residual)\n",
    "      residual = self.norm(name='norm_proj')(residual)\n",
    "\n",
    "    return self.act(residual + y)\n",
    "\n",
    "\n",
    "class _BottleneckResNetBlock(nn.Module):\n",
    "  \"\"\"Bottleneck ResNet block.\"\"\"\n",
    "  filters: int\n",
    "  conv: _ModuleDef\n",
    "  norm: _ModuleDef\n",
    "  act: Callable\n",
    "  strides: tuple[int, int] = (1, 1)\n",
    "\n",
    "  @nn.compact\n",
    "  def __call__(self, x):\n",
    "    residual = x\n",
    "    y = self.conv(self.filters, (1, 1))(x)\n",
    "    y = self.norm()(y)\n",
    "    y = self.act(y)\n",
    "    y = self.conv(self.filters, (3, 3), self.strides)(y)\n",
    "    y = self.norm()(y)\n",
    "    y = self.act(y)\n",
    "    y = self.conv(self.filters * 4, (1, 1))(y)\n",
    "    y = self.norm(scale_init=nn.initializers.zeros)(y)\n",
    "\n",
    "    if residual.shape != y.shape:\n",
    "      residual = self.conv(self.filters * 4, (1, 1),\n",
    "                           self.strides, name='conv_proj')(residual)\n",
    "      residual = self.norm(name='norm_proj')(residual)\n",
    "\n",
    "    return self.act(residual + y)\n",
    "\n",
    "\n",
    "class _ResNet(nn.Module):\n",
    "  \"\"\"ResNetV1.\"\"\"\n",
    "  stage_sizes: Sequence[int]\n",
    "  block_cls: _ModuleDef\n",
    "  num_classes: int\n",
    "  num_filters: int = 64\n",
    "  dtype: Any = jnp.float32\n",
    "  act: Callable = nn.relu\n",
    "  conv: _ModuleDef = nn.Conv\n",
    "\n",
    "  @nn.compact\n",
    "  def __call__(self, x, train: bool = True):\n",
    "    conv = partial(self.conv, use_bias=False, dtype=self.dtype)\n",
    "    norm = partial(nn.BatchNorm,\n",
    "                   use_running_average=not train,\n",
    "                   momentum=0.9,\n",
    "                   epsilon=1e-5,\n",
    "                   dtype=self.dtype)\n",
    "\n",
    "    x = conv(self.num_filters, (7, 7), (2, 2),\n",
    "             padding=[(3, 3), (3, 3)],\n",
    "             name='conv_init')(x)\n",
    "    x = norm(name='bn_init')(x)\n",
    "    x = nn.relu(x)\n",
    "    x = nn.max_pool(x, (3, 3), strides=(2, 2), padding='SAME')\n",
    "    for i, block_size in enumerate(self.stage_sizes):\n",
    "      for j in range(block_size):\n",
    "        strides = (2, 2) if i > 0 and j == 0 else (1, 1)\n",
    "        x = self.block_cls(self.num_filters * 2 ** i,\n",
    "                           strides=strides,\n",
    "                           conv=conv,\n",
    "                           norm=norm,\n",
    "                           act=self.act)(x)\n",
    "    x = jnp.mean(x, axis=(1, 2))\n",
    "    x = nn.Dense(self.num_classes, dtype=self.dtype)(x)\n",
    "    x = jnp.asarray(x, self.dtype)\n",
    "    return x\n",
    "\n",
    "\n",
    "_ResNet18 = partial(_ResNet, stage_sizes=[2, 2, 2, 2],\n",
    "                    block_cls=_ResNetBlock)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "lPh5LGz9JBK_"
   },
   "outputs": [],
   "source": [
    "def get_ntk_fns(O: int):\n",
    "  # Define a ResNet18.\n",
    "  model = _ResNet18(num_classes=O)\n",
    "\n",
    "  # f(x, \\theta)\n",
    "  def apply_fn(params, x):\n",
    "    return model.apply(params, x, train=False, mutable=['batch_stats'])[0]\n",
    "\n",
    "  kwargs = dict(\n",
    "      f=apply_fn,\n",
    "      trace_axes=(),\n",
    "      vmap_axes=0\n",
    "  )\n",
    "\n",
    "  # Different NTK implementations\n",
    "  jacobian_contraction = jit(nt.empirical_ntk_fn(\n",
    "      **kwargs, implementation=nt.NtkImplementation.JACOBIAN_CONTRACTION))\n",
    "  ntvp = jit(nt.empirical_ntk_fn(\n",
    "      **kwargs, implementation=nt.NtkImplementation.NTK_VECTOR_PRODUCTS))\n",
    "  str_derivatives = jit(nt.empirical_ntk_fn(\n",
    "      **kwargs, implementation=nt.NtkImplementation.STRUCTURED_DERIVATIVES))\n",
    "  auto = jit(nt.empirical_ntk_fn(\n",
    "      **kwargs, implementation=nt.NtkImplementation.AUTO))\n",
    "\n",
    "  # Parameters \\theta\n",
    "  params = model.init(random.PRNGKey(0), x1)\n",
    "  return params, (jacobian_contraction, ntvp, str_derivatives, auto)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "0lWFC3QEgao4"
   },
   "source": [
    "# $\\color{blue}O = 8$ logit, batch size $\\color{red}N = 8$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "kbExWKkUg9ew"
   },
   "source": [
    "Structured derivatives compute NTK fastest. NTK-vector products are actually slower in this setting, due to costly forward pass relative to parameters size, and therefore scales poorly with batch size $\\color{red}N$. While it scales better with $\\color{blue}O$ than other methods, it's not enough to overcome the $\\color{red}N^2$ forward passes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "bwhIZWqxKTlt"
   },
   "outputs": [],
   "source": [
    "O = 8\n",
    "N = 8\n",
    "\n",
    "# Input images x\n",
    "input_shape = (224, 224, 3)\n",
    "k1, k2 = random.split(random.PRNGKey(1), 2)\n",
    "x1 = random.normal(k1, (N, *input_shape))\n",
    "x2 = random.normal(k2, (N, *input_shape))\n",
    "\n",
    "params, (ntk_fn_jacobian_contraction, ntk_fn_ntvp, ntk_fn_str_derivatives, ntk_fn_auto) = get_ntk_fns(O=O)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "zObT8WnPggFo",
    "outputId": "0f4ede19-70fb-4d02-f075-b93f2b45d4f8",
    "colab": {
     "base_uri": "https://localhost:8080/"
    }
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "(8, 8, 8, 8)\n"
     ]
    }
   ],
   "source": [
    "# test {\"skip\": true}\n",
    "# Jacobian contraction\n",
    "k_1 = ntk_fn_jacobian_contraction(x1, x2, params)\n",
    "print(k_1.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "FW9gJJ4qggFp",
    "outputId": "37332d1a-7745-4e38-8066-6fbda0ac41e6",
    "colab": {
     "base_uri": "https://localhost:8080/"
    }
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "(8, 8, 8, 8)\n"
     ]
    }
   ],
   "source": [
    "# test {\"skip\": true}\n",
    "# NTK-vector products\n",
    "k_2 = ntk_fn_ntvp(x1, x2, params)\n",
    "print(k_2.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "gFeWnqGQggFp",
    "outputId": "53d93afd-3fbc-49f2-c349-48816c67b92f",
    "colab": {
     "base_uri": "https://localhost:8080/"
    }
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "(8, 8, 8, 8)\n"
     ]
    }
   ],
   "source": [
    "# test {\"skip\": true}\n",
    "# Structured derivatives\n",
    "k_3 = ntk_fn_str_derivatives(x1, x2, params)\n",
    "print(k_3.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Q63v1L1aggFp",
    "outputId": "a430aae2-3a6d-4999-d8b4-c09ecca8b122",
    "colab": {
     "base_uri": "https://localhost:8080/"
    }
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "3.4810557e-06 3.4810557e-06 3.916188e-06\n"
     ]
    }
   ],
   "source": [
    "# test {\"skip\": true}\n",
    "# Make sure kernels agree.\n",
    "print(\n",
    "    jnp.max(jnp.abs(k_1 - k_2)) / jnp.mean(jnp.abs(k_1)),\n",
    "    jnp.max(jnp.abs(k_1 - k_3)) / jnp.mean(jnp.abs(k_1)),\n",
    "    jnp.max(jnp.abs(k_2 - k_3)) / jnp.mean(jnp.abs(k_2))\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Ux7AEZ9fggFp",
    "outputId": "e6c9cedc-1437-4c46-a871-b759d2a005a8",
    "colab": {
     "base_uri": "https://localhost:8080/"
    }
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "impl=1, flops=3964546560.0\n",
      "impl=2, flops=20214009856.0\n",
      "impl=3, flops=4047975424.0\n",
      "(8, 8, 8, 8)\n"
     ]
    }
   ],
   "source": [
    "# Selects best method based on FLOPs at first call / compilation.\n",
    "# Takes about 3x more time to compile.\n",
    "# WARNING: due to an XLA issue, currently only works correctly on TPUs!\n",
    "# Wrong FLOPs for CPU/GPU of JITted functions.\n",
    "k_0 = ntk_fn_auto(x1, x2, params)\n",
    "print(k_0.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "diP7nkBuggFp",
    "outputId": "22189369-9118-466a-991d-f62ba0cabcda",
    "colab": {
     "base_uri": "https://localhost:8080/"
    }
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "1 loop, best of 5: 222 ms per loop\n"
     ]
    }
   ],
   "source": [
    "# test {\"skip\": true}\n",
    "%%timeit\n",
    "ntk_fn_jacobian_contraction(x1, x2, params).block_until_ready()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "wehCdvi2ggFp",
    "outputId": "ae60ad1e-c8cc-4d2b-a650-233c79713783",
    "colab": {
     "base_uri": "https://localhost:8080/"
    }
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "1 loop, best of 5: 311 ms per loop\n"
     ]
    }
   ],
   "source": [
    "# test {\"skip\": true}\n",
    "%%timeit\n",
    "# Slower - forward pass (FP) is expensive relative to parameters.\n",
    "# Time cost scales poorly with batch size N.\n",
    "ntk_fn_ntvp(x1, x2, params).block_until_ready()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Yrm53akVggFp",
    "outputId": "a11a145f-99f4-463f-e73e-95d0e03d2e20",
    "colab": {
     "base_uri": "https://localhost:8080/"
    }
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "10 loops, best of 5: 90.6 ms per loop\n"
     ]
    }
   ],
   "source": [
    "# test {\"skip\": true}\n",
    "%%timeit\n",
    "# 2X faster.\n",
    "ntk_fn_str_derivatives(x1, x2, params).block_until_ready()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "P1QtkBqLggFp",
    "outputId": "d7acf2c8-0d3d-439e-8e1e-bfb3afa31260",
    "colab": {
     "base_uri": "https://localhost:8080/"
    }
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "1 loop, best of 5: 222 ms per loop\n"
     ]
    }
   ],
   "source": [
    "# test {\"skip\": true}\n",
    "%%timeit\n",
    "# On TPU should match the fastest method.\n",
    "# On GPU/CPU, currently is broken, and may not be the fastest.\n",
    "ntk_fn_auto(x1, x2, params).block_until_ready()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "t1DLRESpXg7L"
   },
   "source": [
    "# $\\color{blue}O = 128$ logits, batch size $\\color{red}N = 1$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "dfoHOoT-gGy6"
   },
   "source": [
    "Both NTK-vector products and Structured derivatives compute NTK faster than Jacobian contraction. NTK-vector products incur no penalty when batch size $\\color{red}N = 1$, and leverage their beneficial scaling with large $\\color{blue}O = 128$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "3oHBaAmDhBON",
    "cellView": "code"
   },
   "outputs": [],
   "source": [
    "# test {\"skip\": true}\n",
    "O = 128\n",
    "N = 1\n",
    "\n",
    "# Input images x\n",
    "input_shape = (224, 224, 3)\n",
    "k1, k2 = random.split(random.PRNGKey(1), 2)\n",
    "x1 = random.normal(k1, (N, *input_shape))\n",
    "x2 = random.normal(k2, (N, *input_shape))\n",
    "\n",
    "params, (ntk_fn_jacobian_contraction, ntk_fn_ntvp, ntk_fn_str_derivatives, ntk_fn_auto) = get_ntk_fns(O=O)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "sc1bUvL-KrK9",
    "outputId": "2728b913-394a-40b0-ee1d-ca9b3d8e0ab0",
    "cellView": "code",
    "colab": {
     "base_uri": "https://localhost:8080/"
    }
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "(1, 1, 128, 128)\n"
     ]
    }
   ],
   "source": [
    "# test {\"skip\": true}\n",
    "# Jacobian contraction\n",
    "k_1 = ntk_fn_jacobian_contraction(x1, x2, params)\n",
    "print(k_1.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "qdNsmnjOKyp0",
    "outputId": "a88bf042-9f82-468d-dc71-87b32f92514a",
    "cellView": "code",
    "colab": {
     "base_uri": "https://localhost:8080/"
    }
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "(1, 1, 128, 128)\n"
     ]
    }
   ],
   "source": [
    "# test {\"skip\": true}\n",
    "# NTK-vector products\n",
    "k_2 = ntk_fn_ntvp(x1, x2, params)\n",
    "print(k_2.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Iw6HL260K26E",
    "outputId": "20a90cb0-136a-445a-9085-180f1786e533",
    "colab": {
     "base_uri": "https://localhost:8080/"
    }
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "(1, 1, 128, 128)\n"
     ]
    }
   ],
   "source": [
    "# test {\"skip\": true}\n",
    "# Structured derivatives\n",
    "k_3 = ntk_fn_str_derivatives(x1, x2, params)\n",
    "print(k_3.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "DYG0fV9nOjnd",
    "outputId": "3ac3fc29-1fd2-4688-8793-52c20b345893",
    "cellView": "code",
    "colab": {
     "base_uri": "https://localhost:8080/"
    }
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "1.637553e-05 1.0234707e-05 1.2281647e-05\n"
     ]
    }
   ],
   "source": [
    "# test {\"skip\": true}\n",
    "# Make sure kernels agree.\n",
    "print(\n",
    "    jnp.max(jnp.abs(k_1 - k_2)) / jnp.mean(jnp.abs(k_1)),\n",
    "    jnp.max(jnp.abs(k_1 - k_3)) / jnp.mean(jnp.abs(k_1)),\n",
    "    jnp.max(jnp.abs(k_2 - k_3)) / jnp.mean(jnp.abs(k_2))\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "DyF4M5_HK5Fk",
    "outputId": "8fa15161-df02-4253-8ff3-9a030b3dda56",
    "colab": {
     "base_uri": "https://localhost:8080/"
    }
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "impl=1, flops=6864798208.0\n",
      "impl=2, flops=7510242816.0\n",
      "impl=3, flops=6848109568.0\n",
      "(1, 1, 128, 128)\n"
     ]
    }
   ],
   "source": [
    "# test {\"skip\": true}\n",
    "# Selects best method based on FLOPs at first call / compilation.\n",
    "# Takes about 3x more time to compile.\n",
    "# WARNING: due to an XLA issue, currently only works correctly on TPUs!\n",
    "# Wrong FLOPs for CPU/GPU of JITted functions.\n",
    "k_0 = ntk_fn_auto(x1, x2, params)\n",
    "print(k_0.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "8g1IO71LLJlG",
    "outputId": "6a958f53-67f5-46cc-f502-7dea33498632",
    "cellView": "code",
    "colab": {
     "base_uri": "https://localhost:8080/"
    }
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "1 loop, best of 5: 454 ms per loop\n"
     ]
    }
   ],
   "source": [
    "# test {\"skip\": true}\n",
    "%%timeit\n",
    "ntk_fn_jacobian_contraction(x1, x2, params).block_until_ready()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "KEIPcXYRMys2",
    "outputId": "53bdcc98-2d89-4653-db83-3ecb64f59b75",
    "cellView": "code",
    "colab": {
     "base_uri": "https://localhost:8080/"
    }
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "10 loops, best of 5: 150 ms per loop\n"
     ]
    }
   ],
   "source": [
    "# test {\"skip\": true}\n",
    "%%timeit\n",
    "# 3X faster!\n",
    "ntk_fn_ntvp(x1, x2, params).block_until_ready()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "gVyUPA8xM1ot",
    "outputId": "45ad81e0-d0d5-40cb-9e71-8aaa0370ef09",
    "cellView": "code",
    "colab": {
     "base_uri": "https://localhost:8080/"
    }
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "10 loops, best of 5: 113 ms per loop\n"
     ]
    }
   ],
   "source": [
    "# test {\"skip\": true}\n",
    "%%timeit\n",
    "# 4X faster!\n",
    "ntk_fn_str_derivatives(x1, x2, params).block_until_ready()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "o81r3UHJM34_",
    "outputId": "22947256-12b2-490f-97e1-861ffedadea9",
    "cellView": "code",
    "colab": {
     "base_uri": "https://localhost:8080/"
    }
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "10 loops, best of 5: 113 ms per loop\n"
     ]
    }
   ],
   "source": [
    "# test {\"skip\": true}\n",
    "%%timeit\n",
    "# On TPU should match the fastest method.\n",
    "# On GPU/CPU, currently is broken, and may not be the fastest.\n",
    "ntk_fn_auto(x1, x2, params).block_until_ready()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Dh_aFUZFXm_C"
   },
   "source": [
    "# $\\color{blue}O = 1000$ logits, batch size $\\color{red}N = 1$, full NTK"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "8rs4KF07fjjv"
   },
   "source": [
    "Structured derivatives allows to compute full $1000\\times 1000$ ImageNet NTK. Other methods run out of memory."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Z4MpdSJ7hFsD",
    "cellView": "code"
   },
   "outputs": [],
   "source": [
    "# test {\"skip\": true}\n",
    "O = 1000\n",
    "N = 1\n",
    "\n",
    "# Input images x\n",
    "input_shape = (224, 224, 3)\n",
    "k1, k2 = random.split(random.PRNGKey(1), 2)\n",
    "x1 = random.normal(k1, (N, *input_shape))\n",
    "x2 = random.normal(k2, (N, *input_shape))\n",
    "\n",
    "params, (ntk_fn_jacobian_contraction, ntk_fn_ntvp, ntk_fn_str_derivatives, ntk_fn_auto) = get_ntk_fns(O=O)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "ynBEtkA7d_2G",
    "outputId": "a7f6d320-530e-475e-a4a5-15fb3a2e6efd",
    "cellView": "code",
    "colab": {
     "base_uri": "https://localhost:8080/"
    }
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "(1, 1, 1000, 1000)\n"
     ]
    }
   ],
   "source": [
    "# test {\"skip\": true}\n",
    "# Structured derivatives - fits in memory!\n",
    "k_3 = ntk_fn_str_derivatives(x1, x2, params)\n",
    "print(k_3.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Afniv9k2d_2H",
    "outputId": "01bec9ff-af1b-4dde-c517-0c06c99452b4",
    "cellView": "code",
    "colab": {
     "base_uri": "https://localhost:8080/"
    }
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "1 loop, best of 5: 985 ms per loop\n"
     ]
    }
   ],
   "source": [
    "# test {\"skip\": true}\n",
    "%%timeit\n",
    "ntk_fn_str_derivatives(x1, x2, params).block_until_ready()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "height": 867,
     "base_uri": "https://localhost:8080/"
    },
    "id": "4Ou5a6y4zH5Q",
    "outputId": "3f558f91-9f22-4bf9-86f4-e76a15835121",
    "cellView": "code"
   },
   "outputs": [
    {
     "output_type": "error",
     "ename": "XlaRuntimeError",
     "evalue": "ignored",
     "traceback": [
      "\u001B[0;31m---------------------------------------------------------------------------\u001B[0m",
      "\u001B[0;31mUnfilteredStackTrace\u001B[0m                      Traceback (most recent call last)",
      "\u001B[0;32m<ipython-input-29-43ee9476f982>\u001B[0m in \u001B[0;36m<module>\u001B[0;34m()\u001B[0m\n\u001B[1;32m      1\u001B[0m \u001B[0;31m# NTK-vector products - OOM!\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[0;32m----> 2\u001B[0;31m \u001B[0mk_3\u001B[0m \u001B[0;34m=\u001B[0m \u001B[0mntk_fn_ntvp\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0mx1\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0mx2\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0mparams\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[0m\u001B[1;32m      3\u001B[0m \u001B[0mprint\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0mk_3\u001B[0m\u001B[0;34m.\u001B[0m\u001B[0mshape\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n",
      "\u001B[0;32m/usr/local/lib/python3.7/dist-packages/jax/_src/traceback_util.py\u001B[0m in \u001B[0;36mreraise_with_filtered_traceback\u001B[0;34m(*args, **kwargs)\u001B[0m\n\u001B[1;32m    161\u001B[0m     \u001B[0;32mtry\u001B[0m\u001B[0;34m:\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[0;32m--> 162\u001B[0;31m       \u001B[0;32mreturn\u001B[0m \u001B[0mfun\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0;34m*\u001B[0m\u001B[0margs\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0;34m**\u001B[0m\u001B[0mkwargs\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[0m\u001B[1;32m    163\u001B[0m     \u001B[0;32mexcept\u001B[0m \u001B[0mException\u001B[0m \u001B[0;32mas\u001B[0m \u001B[0me\u001B[0m\u001B[0;34m:\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n",
      "\u001B[0;32m/usr/local/lib/python3.7/dist-packages/jax/_src/api.py\u001B[0m in \u001B[0;36mcache_miss\u001B[0;34m(*args, **kwargs)\u001B[0m\n\u001B[1;32m    475\u001B[0m         \u001B[0mdevice\u001B[0m\u001B[0;34m=\u001B[0m\u001B[0mdevice\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0mbackend\u001B[0m\u001B[0;34m=\u001B[0m\u001B[0mbackend\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0mname\u001B[0m\u001B[0;34m=\u001B[0m\u001B[0mflat_fun\u001B[0m\u001B[0;34m.\u001B[0m\u001B[0m__name__\u001B[0m\u001B[0;34m,\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[0;32m--> 476\u001B[0;31m         donated_invars=donated_invars, inline=inline, keep_unused=keep_unused)\n\u001B[0m\u001B[1;32m    477\u001B[0m     \u001B[0mout_pytree_def\u001B[0m \u001B[0;34m=\u001B[0m \u001B[0mout_tree\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n",
      "\u001B[0;32m/usr/local/lib/python3.7/dist-packages/jax/core.py\u001B[0m in \u001B[0;36mbind\u001B[0;34m(self, fun, *args, **params)\u001B[0m\n\u001B[1;32m   1764\u001B[0m   \u001B[0;32mdef\u001B[0m \u001B[0mbind\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0mself\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0mfun\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0;34m*\u001B[0m\u001B[0margs\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0;34m**\u001B[0m\u001B[0mparams\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m:\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[0;32m-> 1765\u001B[0;31m     \u001B[0;32mreturn\u001B[0m \u001B[0mcall_bind\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0mself\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0mfun\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0;34m*\u001B[0m\u001B[0margs\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0;34m**\u001B[0m\u001B[0mparams\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[0m\u001B[1;32m   1766\u001B[0m \u001B[0;34m\u001B[0m\u001B[0m\n",
      "\u001B[0;32m/usr/local/lib/python3.7/dist-packages/jax/core.py\u001B[0m in \u001B[0;36mcall_bind\u001B[0;34m(primitive, fun, *args, **params)\u001B[0m\n\u001B[1;32m   1780\u001B[0m   \u001B[0mfun_\u001B[0m \u001B[0;34m=\u001B[0m \u001B[0mlu\u001B[0m\u001B[0;34m.\u001B[0m\u001B[0mannotate\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0mfun_\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0mfun\u001B[0m\u001B[0;34m.\u001B[0m\u001B[0min_type\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[0;32m-> 1781\u001B[0;31m   \u001B[0mouts\u001B[0m \u001B[0;34m=\u001B[0m \u001B[0mtop_trace\u001B[0m\u001B[0;34m.\u001B[0m\u001B[0mprocess_call\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0mprimitive\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0mfun_\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0mtracers\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0mparams\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[0m\u001B[1;32m   1782\u001B[0m   \u001B[0;32mreturn\u001B[0m \u001B[0mmap\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0mfull_lower\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0mapply_todos\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0menv_trace_todo\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0mouts\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n",
      "\u001B[0;32m/usr/local/lib/python3.7/dist-packages/jax/core.py\u001B[0m in \u001B[0;36mprocess_call\u001B[0;34m(self, primitive, f, tracers, params)\u001B[0m\n\u001B[1;32m    677\u001B[0m   \u001B[0;32mdef\u001B[0m \u001B[0mprocess_call\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0mself\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0mprimitive\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0mf\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0mtracers\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0mparams\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m:\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[0;32m--> 678\u001B[0;31m     \u001B[0;32mreturn\u001B[0m \u001B[0mprimitive\u001B[0m\u001B[0;34m.\u001B[0m\u001B[0mimpl\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0mf\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0;34m*\u001B[0m\u001B[0mtracers\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0;34m**\u001B[0m\u001B[0mparams\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[0m\u001B[1;32m    679\u001B[0m   \u001B[0mprocess_map\u001B[0m \u001B[0;34m=\u001B[0m \u001B[0mprocess_call\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n",
      "\u001B[0;32m/usr/local/lib/python3.7/dist-packages/jax/_src/dispatch.py\u001B[0m in \u001B[0;36m_xla_call_impl\u001B[0;34m(***failed resolving arguments***)\u001B[0m\n\u001B[1;32m    182\u001B[0m   compiled_fun = _xla_callable(fun, device, backend, name, donated_invars,\n\u001B[0;32m--> 183\u001B[0;31m                                keep_unused, *arg_specs)\n\u001B[0m\u001B[1;32m    184\u001B[0m   \u001B[0;32mtry\u001B[0m\u001B[0;34m:\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n",
      "\u001B[0;32m/usr/local/lib/python3.7/dist-packages/jax/linear_util.py\u001B[0m in \u001B[0;36mmemoized_fun\u001B[0;34m(fun, *args)\u001B[0m\n\u001B[1;32m    284\u001B[0m     \u001B[0;32melse\u001B[0m\u001B[0;34m:\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[0;32m--> 285\u001B[0;31m       \u001B[0mans\u001B[0m \u001B[0;34m=\u001B[0m \u001B[0mcall\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0mfun\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0;34m*\u001B[0m\u001B[0margs\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[0m\u001B[1;32m    286\u001B[0m       \u001B[0mcache\u001B[0m\u001B[0;34m[\u001B[0m\u001B[0mkey\u001B[0m\u001B[0;34m]\u001B[0m \u001B[0;34m=\u001B[0m \u001B[0;34m(\u001B[0m\u001B[0mans\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0mfun\u001B[0m\u001B[0;34m.\u001B[0m\u001B[0mstores\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n",
      "\u001B[0;32m/usr/local/lib/python3.7/dist-packages/jax/_src/dispatch.py\u001B[0m in \u001B[0;36m_xla_callable_uncached\u001B[0;34m(fun, device, backend, name, donated_invars, keep_unused, *arg_specs)\u001B[0m\n\u001B[1;32m    230\u001B[0m   return lower_xla_callable(fun, device, backend, name, donated_invars, False,\n\u001B[0;32m--> 231\u001B[0;31m                             keep_unused, *arg_specs).compile().unsafe_call\n\u001B[0m\u001B[1;32m    232\u001B[0m \u001B[0;34m\u001B[0m\u001B[0m\n",
      "\u001B[0;32m/usr/local/lib/python3.7/dist-packages/jax/_src/dispatch.py\u001B[0m in \u001B[0;36mcompile\u001B[0;34m(self)\u001B[0m\n\u001B[1;32m    704\u001B[0m         self._executable = XlaCompiledComputation.from_xla_computation(\n\u001B[0;32m--> 705\u001B[0;31m             self.name, self._hlo, self._explicit_args, **self.compile_args)\n\u001B[0m\u001B[1;32m    706\u001B[0m \u001B[0;34m\u001B[0m\u001B[0m\n",
      "\u001B[0;32m/usr/local/lib/python3.7/dist-packages/jax/_src/dispatch.py\u001B[0m in \u001B[0;36mfrom_xla_computation\u001B[0;34m(name, xla_computation, explicit_args, nreps, device, backend, tuple_args, in_avals, out_avals, effects, kept_var_idx, keepalive)\u001B[0m\n\u001B[1;32m    805\u001B[0m                           \"in {elapsed_time} sec\"):\n\u001B[0;32m--> 806\u001B[0;31m       \u001B[0mcompiled\u001B[0m \u001B[0;34m=\u001B[0m \u001B[0mcompile_or_get_cached\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0mbackend\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0mxla_computation\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0moptions\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[0m\u001B[1;32m    807\u001B[0m     buffer_counts = (None if len(out_avals) == 1 and not config.jax_dynamic_shapes\n",
      "\u001B[0;32m/usr/local/lib/python3.7/dist-packages/jax/_src/dispatch.py\u001B[0m in \u001B[0;36mcompile_or_get_cached\u001B[0;34m(backend, computation, compile_options)\u001B[0m\n\u001B[1;32m    767\u001B[0m     \u001B[0m_dump_ir_to_file\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0mmodule_name\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0mir_str\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[0;32m--> 768\u001B[0;31m   \u001B[0;32mreturn\u001B[0m \u001B[0mbackend_compile\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0mbackend\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0mcomputation\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0mcompile_options\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[0m\u001B[1;32m    769\u001B[0m \u001B[0;34m\u001B[0m\u001B[0m\n",
      "\u001B[0;32m/usr/local/lib/python3.7/dist-packages/jax/_src/profiler.py\u001B[0m in \u001B[0;36mwrapper\u001B[0;34m(*args, **kwargs)\u001B[0m\n\u001B[1;32m    205\u001B[0m     \u001B[0;32mwith\u001B[0m \u001B[0mTraceAnnotation\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0mname\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0;34m**\u001B[0m\u001B[0mdecorator_kwargs\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m:\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[0;32m--> 206\u001B[0;31m       \u001B[0;32mreturn\u001B[0m \u001B[0mfunc\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0;34m*\u001B[0m\u001B[0margs\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0;34m**\u001B[0m\u001B[0mkwargs\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[0m\u001B[1;32m    207\u001B[0m     \u001B[0;32mreturn\u001B[0m \u001B[0mwrapper\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n",
      "\u001B[0;32m/usr/local/lib/python3.7/dist-packages/jax/_src/dispatch.py\u001B[0m in \u001B[0;36mbackend_compile\u001B[0;34m(backend, built_c, options)\u001B[0m\n\u001B[1;32m    712\u001B[0m   \u001B[0;31m# separately in Python profiling results\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[0;32m--> 713\u001B[0;31m   \u001B[0;32mreturn\u001B[0m \u001B[0mbackend\u001B[0m\u001B[0;34m.\u001B[0m\u001B[0mcompile\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0mbuilt_c\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0mcompile_options\u001B[0m\u001B[0;34m=\u001B[0m\u001B[0moptions\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[0m\u001B[1;32m    714\u001B[0m \u001B[0;34m\u001B[0m\u001B[0m\n",
      "\u001B[0;31mUnfilteredStackTrace\u001B[0m: jaxlib.xla_extension.XlaRuntimeError: UNKNOWN: Failed to determine best cudnn convolution algorithm for:\n%cudnn-conv-bw-filter = (f32[3,3,512,512000]{1,0,2,3}, u8[0]{0}) custom-call(f32[1,7,7,512]{2,1,3,0} %bitcast.332, f32[1,7,7,512000]{2,1,3,0} %bitcast.839), window={size=3x3 pad=1_1x1_1}, dim_labels=b01f_01io->b01f, custom_call_target=\"__cudnn$convBackwardFilter\", metadata={op_name=\"jit(ntk_fn)/jit(main)/conv_general_dilated[window_strides=(1, 1) padding=((1, 1), (1, 1)) lhs_dilation=(1, 1) rhs_dilation=(1, 1) dimension_numbers=ConvDimensionNumbers(lhs_spec=(3, 0, 1, 2), rhs_spec=(3, 0, 1, 2), out_spec=(2, 3, 0, 1)) feature_group_count=1 batch_group_count=1 lhs_shape=(1, 7, 7, 512) rhs_shape=(1, 7, 7, 512000) precision=None preferred_element_type=None]\" source_file=\"/usr/local/lib/python3.7/dist-packages/flax/linen/linear.py\" source_line=434}, backend_config=\"{\\\"conv_result_scale\\\":1,\\\"activation_mode\\\":\\\"0\\\",\\\"side_input_scale\\\":0}\"\n\nOriginal error: INTERNAL: stream did not block host until done; was already in an error state\n\nTo ignore this failure and try to use a fallback algorithm (which may have suboptimal performance), use XLA_FLAGS=--xla_gpu_strict_conv_algorithm_picker=false.  Please also file a bug for the root cause of failing autotuning.\n\nThe stack trace below excludes JAX-internal frames.\nThe preceding is the original exception that occurred, unmodified.\n\n--------------------",
      "\nThe above exception was the direct cause of the following exception:\n",
      "\u001B[0;31mXlaRuntimeError\u001B[0m                           Traceback (most recent call last)",
      "\u001B[0;32m<ipython-input-29-43ee9476f982>\u001B[0m in \u001B[0;36m<module>\u001B[0;34m()\u001B[0m\n\u001B[1;32m      1\u001B[0m \u001B[0;31m# NTK-vector products - OOM!\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[0;32m----> 2\u001B[0;31m \u001B[0mk_3\u001B[0m \u001B[0;34m=\u001B[0m \u001B[0mntk_fn_ntvp\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0mx1\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0mx2\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0mparams\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[0m\u001B[1;32m      3\u001B[0m \u001B[0mprint\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0mk_3\u001B[0m\u001B[0;34m.\u001B[0m\u001B[0mshape\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n",
      "\u001B[0;31mXlaRuntimeError\u001B[0m: UNKNOWN: Failed to determine best cudnn convolution algorithm for:\n%cudnn-conv-bw-filter = (f32[3,3,512,512000]{1,0,2,3}, u8[0]{0}) custom-call(f32[1,7,7,512]{2,1,3,0} %bitcast.332, f32[1,7,7,512000]{2,1,3,0} %bitcast.839), window={size=3x3 pad=1_1x1_1}, dim_labels=b01f_01io->b01f, custom_call_target=\"__cudnn$convBackwardFilter\", metadata={op_name=\"jit(ntk_fn)/jit(main)/conv_general_dilated[window_strides=(1, 1) padding=((1, 1), (1, 1)) lhs_dilation=(1, 1) rhs_dilation=(1, 1) dimension_numbers=ConvDimensionNumbers(lhs_spec=(3, 0, 1, 2), rhs_spec=(3, 0, 1, 2), out_spec=(2, 3, 0, 1)) feature_group_count=1 batch_group_count=1 lhs_shape=(1, 7, 7, 512) rhs_shape=(1, 7, 7, 512000) precision=None preferred_element_type=None]\" source_file=\"/usr/local/lib/python3.7/dist-packages/flax/linen/linear.py\" source_line=434}, backend_config=\"{\\\"conv_result_scale\\\":1,\\\"activation_mode\\\":\\\"0\\\",\\\"side_input_scale\\\":0}\"\n\nOriginal error: INTERNAL: stream did not block host until done; was already in an error state\n\nTo ignore this failure and try to use a fallback algorithm (which may have suboptimal performance), use XLA_FLAGS=--xla_gpu_strict_conv_algorithm_picker=false.  Please also file a bug for the root cause of failing autotuning."
     ]
    }
   ],
   "source": [
    "# test {\"skip\": true}\n",
    "# NTK-vector products - OOM!\n",
    "k_3 = ntk_fn_ntvp(x1, x2, params)\n",
    "print(k_3.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "height": 936,
     "base_uri": "https://localhost:8080/"
    },
    "id": "rjXRalQfd_2G",
    "outputId": "e442c5ea-af41-462b-a4bb-266855ccb3c7",
    "cellView": "code"
   },
   "outputs": [
    {
     "output_type": "error",
     "ename": "XlaRuntimeError",
     "evalue": "ignored",
     "traceback": [
      "\u001B[0;31m---------------------------------------------------------------------------\u001B[0m",
      "\u001B[0;31mUnfilteredStackTrace\u001B[0m                      Traceback (most recent call last)",
      "\u001B[0;32m<ipython-input-30-ae557f05f951>\u001B[0m in \u001B[0;36m<module>\u001B[0;34m()\u001B[0m\n\u001B[1;32m      1\u001B[0m \u001B[0;31m# Jacobian contraction - OOM!\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[0;32m----> 2\u001B[0;31m \u001B[0mk_1\u001B[0m \u001B[0;34m=\u001B[0m \u001B[0mntk_fn_jacobian_contraction\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0mx1\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0mx2\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0mparams\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[0m\u001B[1;32m      3\u001B[0m \u001B[0mprint\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0mk_1\u001B[0m\u001B[0;34m.\u001B[0m\u001B[0mshape\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n",
      "\u001B[0;32m/usr/local/lib/python3.7/dist-packages/jax/_src/traceback_util.py\u001B[0m in \u001B[0;36mreraise_with_filtered_traceback\u001B[0;34m(*args, **kwargs)\u001B[0m\n\u001B[1;32m    161\u001B[0m     \u001B[0;32mtry\u001B[0m\u001B[0;34m:\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[0;32m--> 162\u001B[0;31m       \u001B[0;32mreturn\u001B[0m \u001B[0mfun\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0;34m*\u001B[0m\u001B[0margs\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0;34m**\u001B[0m\u001B[0mkwargs\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[0m\u001B[1;32m    163\u001B[0m     \u001B[0;32mexcept\u001B[0m \u001B[0mException\u001B[0m \u001B[0;32mas\u001B[0m \u001B[0me\u001B[0m\u001B[0;34m:\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n",
      "\u001B[0;32m/usr/local/lib/python3.7/dist-packages/jax/_src/api.py\u001B[0m in \u001B[0;36mcache_miss\u001B[0;34m(*args, **kwargs)\u001B[0m\n\u001B[1;32m    475\u001B[0m         \u001B[0mdevice\u001B[0m\u001B[0;34m=\u001B[0m\u001B[0mdevice\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0mbackend\u001B[0m\u001B[0;34m=\u001B[0m\u001B[0mbackend\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0mname\u001B[0m\u001B[0;34m=\u001B[0m\u001B[0mflat_fun\u001B[0m\u001B[0;34m.\u001B[0m\u001B[0m__name__\u001B[0m\u001B[0;34m,\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[0;32m--> 476\u001B[0;31m         donated_invars=donated_invars, inline=inline, keep_unused=keep_unused)\n\u001B[0m\u001B[1;32m    477\u001B[0m     \u001B[0mout_pytree_def\u001B[0m \u001B[0;34m=\u001B[0m \u001B[0mout_tree\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n",
      "\u001B[0;32m/usr/local/lib/python3.7/dist-packages/jax/core.py\u001B[0m in \u001B[0;36mbind\u001B[0;34m(self, fun, *args, **params)\u001B[0m\n\u001B[1;32m   1764\u001B[0m   \u001B[0;32mdef\u001B[0m \u001B[0mbind\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0mself\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0mfun\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0;34m*\u001B[0m\u001B[0margs\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0;34m**\u001B[0m\u001B[0mparams\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m:\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[0;32m-> 1765\u001B[0;31m     \u001B[0;32mreturn\u001B[0m \u001B[0mcall_bind\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0mself\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0mfun\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0;34m*\u001B[0m\u001B[0margs\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0;34m**\u001B[0m\u001B[0mparams\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[0m\u001B[1;32m   1766\u001B[0m \u001B[0;34m\u001B[0m\u001B[0m\n",
      "\u001B[0;32m/usr/local/lib/python3.7/dist-packages/jax/core.py\u001B[0m in \u001B[0;36mcall_bind\u001B[0;34m(primitive, fun, *args, **params)\u001B[0m\n\u001B[1;32m   1780\u001B[0m   \u001B[0mfun_\u001B[0m \u001B[0;34m=\u001B[0m \u001B[0mlu\u001B[0m\u001B[0;34m.\u001B[0m\u001B[0mannotate\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0mfun_\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0mfun\u001B[0m\u001B[0;34m.\u001B[0m\u001B[0min_type\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[0;32m-> 1781\u001B[0;31m   \u001B[0mouts\u001B[0m \u001B[0;34m=\u001B[0m \u001B[0mtop_trace\u001B[0m\u001B[0;34m.\u001B[0m\u001B[0mprocess_call\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0mprimitive\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0mfun_\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0mtracers\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0mparams\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[0m\u001B[1;32m   1782\u001B[0m   \u001B[0;32mreturn\u001B[0m \u001B[0mmap\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0mfull_lower\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0mapply_todos\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0menv_trace_todo\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0mouts\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n",
      "\u001B[0;32m/usr/local/lib/python3.7/dist-packages/jax/core.py\u001B[0m in \u001B[0;36mprocess_call\u001B[0;34m(self, primitive, f, tracers, params)\u001B[0m\n\u001B[1;32m    677\u001B[0m   \u001B[0;32mdef\u001B[0m \u001B[0mprocess_call\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0mself\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0mprimitive\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0mf\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0mtracers\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0mparams\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m:\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[0;32m--> 678\u001B[0;31m     \u001B[0;32mreturn\u001B[0m \u001B[0mprimitive\u001B[0m\u001B[0;34m.\u001B[0m\u001B[0mimpl\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0mf\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0;34m*\u001B[0m\u001B[0mtracers\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0;34m**\u001B[0m\u001B[0mparams\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[0m\u001B[1;32m    679\u001B[0m   \u001B[0mprocess_map\u001B[0m \u001B[0;34m=\u001B[0m \u001B[0mprocess_call\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n",
      "\u001B[0;32m/usr/local/lib/python3.7/dist-packages/jax/_src/dispatch.py\u001B[0m in \u001B[0;36m_xla_call_impl\u001B[0;34m(***failed resolving arguments***)\u001B[0m\n\u001B[1;32m    182\u001B[0m   compiled_fun = _xla_callable(fun, device, backend, name, donated_invars,\n\u001B[0;32m--> 183\u001B[0;31m                                keep_unused, *arg_specs)\n\u001B[0m\u001B[1;32m    184\u001B[0m   \u001B[0;32mtry\u001B[0m\u001B[0;34m:\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n",
      "\u001B[0;32m/usr/local/lib/python3.7/dist-packages/jax/linear_util.py\u001B[0m in \u001B[0;36mmemoized_fun\u001B[0;34m(fun, *args)\u001B[0m\n\u001B[1;32m    284\u001B[0m     \u001B[0;32melse\u001B[0m\u001B[0;34m:\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[0;32m--> 285\u001B[0;31m       \u001B[0mans\u001B[0m \u001B[0;34m=\u001B[0m \u001B[0mcall\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0mfun\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0;34m*\u001B[0m\u001B[0margs\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[0m\u001B[1;32m    286\u001B[0m       \u001B[0mcache\u001B[0m\u001B[0;34m[\u001B[0m\u001B[0mkey\u001B[0m\u001B[0;34m]\u001B[0m \u001B[0;34m=\u001B[0m \u001B[0;34m(\u001B[0m\u001B[0mans\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0mfun\u001B[0m\u001B[0;34m.\u001B[0m\u001B[0mstores\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n",
      "\u001B[0;32m/usr/local/lib/python3.7/dist-packages/jax/_src/dispatch.py\u001B[0m in \u001B[0;36m_xla_callable_uncached\u001B[0;34m(fun, device, backend, name, donated_invars, keep_unused, *arg_specs)\u001B[0m\n\u001B[1;32m    230\u001B[0m   return lower_xla_callable(fun, device, backend, name, donated_invars, False,\n\u001B[0;32m--> 231\u001B[0;31m                             keep_unused, *arg_specs).compile().unsafe_call\n\u001B[0m\u001B[1;32m    232\u001B[0m \u001B[0;34m\u001B[0m\u001B[0m\n",
      "\u001B[0;32m/usr/local/lib/python3.7/dist-packages/jax/_src/dispatch.py\u001B[0m in \u001B[0;36mcompile\u001B[0;34m(self)\u001B[0m\n\u001B[1;32m    704\u001B[0m         self._executable = XlaCompiledComputation.from_xla_computation(\n\u001B[0;32m--> 705\u001B[0;31m             self.name, self._hlo, self._explicit_args, **self.compile_args)\n\u001B[0m\u001B[1;32m    706\u001B[0m \u001B[0;34m\u001B[0m\u001B[0m\n",
      "\u001B[0;32m/usr/local/lib/python3.7/dist-packages/jax/_src/dispatch.py\u001B[0m in \u001B[0;36mfrom_xla_computation\u001B[0;34m(name, xla_computation, explicit_args, nreps, device, backend, tuple_args, in_avals, out_avals, effects, kept_var_idx, keepalive)\u001B[0m\n\u001B[1;32m    805\u001B[0m                           \"in {elapsed_time} sec\"):\n\u001B[0;32m--> 806\u001B[0;31m       \u001B[0mcompiled\u001B[0m \u001B[0;34m=\u001B[0m \u001B[0mcompile_or_get_cached\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0mbackend\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0mxla_computation\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0moptions\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[0m\u001B[1;32m    807\u001B[0m     buffer_counts = (None if len(out_avals) == 1 and not config.jax_dynamic_shapes\n",
      "\u001B[0;32m/usr/local/lib/python3.7/dist-packages/jax/_src/dispatch.py\u001B[0m in \u001B[0;36mcompile_or_get_cached\u001B[0;34m(backend, computation, compile_options)\u001B[0m\n\u001B[1;32m    767\u001B[0m     \u001B[0m_dump_ir_to_file\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0mmodule_name\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0mir_str\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[0;32m--> 768\u001B[0;31m   \u001B[0;32mreturn\u001B[0m \u001B[0mbackend_compile\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0mbackend\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0mcomputation\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0mcompile_options\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[0m\u001B[1;32m    769\u001B[0m \u001B[0;34m\u001B[0m\u001B[0m\n",
      "\u001B[0;32m/usr/local/lib/python3.7/dist-packages/jax/_src/profiler.py\u001B[0m in \u001B[0;36mwrapper\u001B[0;34m(*args, **kwargs)\u001B[0m\n\u001B[1;32m    205\u001B[0m     \u001B[0;32mwith\u001B[0m \u001B[0mTraceAnnotation\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0mname\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0;34m**\u001B[0m\u001B[0mdecorator_kwargs\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m:\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[0;32m--> 206\u001B[0;31m       \u001B[0;32mreturn\u001B[0m \u001B[0mfunc\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0;34m*\u001B[0m\u001B[0margs\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0;34m**\u001B[0m\u001B[0mkwargs\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[0m\u001B[1;32m    207\u001B[0m     \u001B[0;32mreturn\u001B[0m \u001B[0mwrapper\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n",
      "\u001B[0;32m/usr/local/lib/python3.7/dist-packages/jax/_src/dispatch.py\u001B[0m in \u001B[0;36mbackend_compile\u001B[0;34m(backend, built_c, options)\u001B[0m\n\u001B[1;32m    712\u001B[0m   \u001B[0;31m# separately in Python profiling results\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[0;32m--> 713\u001B[0;31m   \u001B[0;32mreturn\u001B[0m \u001B[0mbackend\u001B[0m\u001B[0;34m.\u001B[0m\u001B[0mcompile\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0mbuilt_c\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0mcompile_options\u001B[0m\u001B[0;34m=\u001B[0m\u001B[0moptions\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[0m\u001B[1;32m    714\u001B[0m \u001B[0;34m\u001B[0m\u001B[0m\n",
      "\u001B[0;31mUnfilteredStackTrace\u001B[0m: jaxlib.xla_extension.XlaRuntimeError: UNKNOWN: Failed to determine best cudnn convolution algorithm for:\n%cudnn-conv-bw-filter = (f32[3,3,512,512000]{1,0,2,3}, u8[0]{0}) custom-call(f32[1,7,7,512]{2,1,3,0} %bitcast.566, f32[1,7,7,512000]{2,1,3,0} %bitcast.1618), window={size=3x3 pad=1_1x1_1}, dim_labels=b01f_01io->b01f, custom_call_target=\"__cudnn$convBackwardFilter\", metadata={op_name=\"jit(ntk_fn)/jit(main)/conv_general_dilated[window_strides=(1, 1) padding=((1, 1), (1, 1)) lhs_dilation=(1, 1) rhs_dilation=(1, 1) dimension_numbers=ConvDimensionNumbers(lhs_spec=(3, 0, 1, 2), rhs_spec=(3, 0, 1, 2), out_spec=(2, 3, 0, 1)) feature_group_count=1 batch_group_count=1 lhs_shape=(1, 7, 7, 512) rhs_shape=(1, 7, 7, 512000) precision=None preferred_element_type=None]\" source_file=\"/usr/local/lib/python3.7/dist-packages/flax/linen/linear.py\" source_line=434}, backend_config=\"{\\\"conv_result_scale\\\":1,\\\"activation_mode\\\":\\\"0\\\",\\\"side_input_scale\\\":0}\"\n\nOriginal error: INTERNAL: Failed to synchronize GPU for autotuning conv instruction: (f32[3,3,512,512000]{1,0,2,3}, u8[0]{0}) custom-call(f32[1,7,7,512]{2,1,3,0}, f32[1,7,7,512000]{2,1,3,0}), window={size=3x3 pad=1_1x1_1}, dim_labels=b01f_01io->b01f, custom_call_target=\"__cudnn$convBackwardFilter\", backend_config=\"{\\\"conv_result_scale\\\":1,\\\"activation_mode\\\":\\\"0\\\",\\\"side_input_scale\\\":0}\"\n\nTo ignore this failure and try to use a fallback algorithm (which may have suboptimal performance), use XLA_FLAGS=--xla_gpu_strict_conv_algorithm_picker=false.  Please also file a bug for the root cause of failing autotuning.\n\nThe stack trace below excludes JAX-internal frames.\nThe preceding is the original exception that occurred, unmodified.\n\n--------------------",
      "\nThe above exception was the direct cause of the following exception:\n",
      "\u001B[0;31mXlaRuntimeError\u001B[0m                           Traceback (most recent call last)",
      "\u001B[0;32m<ipython-input-30-ae557f05f951>\u001B[0m in \u001B[0;36m<module>\u001B[0;34m()\u001B[0m\n\u001B[1;32m      1\u001B[0m \u001B[0;31m# Jacobian contraction - OOM!\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[0;32m----> 2\u001B[0;31m \u001B[0mk_1\u001B[0m \u001B[0;34m=\u001B[0m \u001B[0mntk_fn_jacobian_contraction\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0mx1\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0mx2\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0mparams\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[0m\u001B[1;32m      3\u001B[0m \u001B[0mprint\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0mk_1\u001B[0m\u001B[0;34m.\u001B[0m\u001B[0mshape\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n",
      "\u001B[0;31mXlaRuntimeError\u001B[0m: UNKNOWN: Failed to determine best cudnn convolution algorithm for:\n%cudnn-conv-bw-filter = (f32[3,3,512,512000]{1,0,2,3}, u8[0]{0}) custom-call(f32[1,7,7,512]{2,1,3,0} %bitcast.566, f32[1,7,7,512000]{2,1,3,0} %bitcast.1618), window={size=3x3 pad=1_1x1_1}, dim_labels=b01f_01io->b01f, custom_call_target=\"__cudnn$convBackwardFilter\", metadata={op_name=\"jit(ntk_fn)/jit(main)/conv_general_dilated[window_strides=(1, 1) padding=((1, 1), (1, 1)) lhs_dilation=(1, 1) rhs_dilation=(1, 1) dimension_numbers=ConvDimensionNumbers(lhs_spec=(3, 0, 1, 2), rhs_spec=(3, 0, 1, 2), out_spec=(2, 3, 0, 1)) feature_group_count=1 batch_group_count=1 lhs_shape=(1, 7, 7, 512) rhs_shape=(1, 7, 7, 512000) precision=None preferred_element_type=None]\" source_file=\"/usr/local/lib/python3.7/dist-packages/flax/linen/linear.py\" source_line=434}, backend_config=\"{\\\"conv_result_scale\\\":1,\\\"activation_mode\\\":\\\"0\\\",\\\"side_input_scale\\\":0}\"\n\nOriginal error: INTERNAL: Failed to synchronize GPU for autotuning conv instruction: (f32[3,3,512,512000]{1,0,2,3}, u8[0]{0}) custom-call(f32[1,7,7,512]{2,1,3,0}, f32[1,7,7,512000]{2,1,3,0}), window={size=3x3 pad=1_1x1_1}, dim_labels=b01f_01io->b01f, custom_call_target=\"__cudnn$convBackwardFilter\", backend_config=\"{\\\"conv_result_scale\\\":1,\\\"activation_mode\\\":\\\"0\\\",\\\"side_input_scale\\\":0}\"\n\nTo ignore this failure and try to use a fallback algorithm (which may have suboptimal performance), use XLA_FLAGS=--xla_gpu_strict_conv_algorithm_picker=false.  Please also file a bug for the root cause of failing autotuning."
     ]
    }
   ],
   "source": [
    "# test {\"skip\": true}\n",
    "# Jacobian contraction - OOM!\n",
    "k_1 = ntk_fn_jacobian_contraction(x1, x2, params)\n",
    "print(k_1.shape)\n"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "collapsed_sections": [],
   "name": "empirical_ntk_resnet.ipynb",
   "provenance": [],
   "toc_visible": true
  },
  "kernelspec": {
   "display_name": "Python 3",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
